#!/usr/bin/env python3
"""
Find correct layer indices for P3, P4, P5 features in YOLOv9
Run this ONCE to identify the correct indices, then update extract_features_for_stata()
"""

import torch
import sys
from pathlib import Path

# Add yolov9 to path
ROOT = Path(__file__).resolve().parent
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))

from models.yolo import DetectionModel


def find_feature_layers(model, input_size=640):
    """
    Find which layers produce P3, P4, P5 features
    
    P3: 80x80 feature map (stride 8)
    P4: 40x40 feature map (stride 16)
    P5: 20x20 feature map (stride 20)
    """
    print("\n" + "="*80)
    print("Finding P3, P4, P5 Layer Indices for YOLOv9")
    print("="*80)
    
    device = next(model.parameters()).device
    x = torch.randn(1, 3, input_size, input_size).to(device)
    
    y = []
    p3_candidates = []
    p4_candidates = []
    p5_candidates = []
    
    print(f"\nInput: {x.shape}")
    print(f"\nScanning all layers...\n")
    print(f"{'Layer':<6} {'Module':<35} {'Output Shape':<30} {'H':<5} {'W':<5} {'Ch':<6}")
    print("-" * 95)
    
    for i, m in enumerate(model.model):
        # Handle multi-input layers
        if m.f != -1:
            x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]
        
        # Forward pass
        x = m(x)
        
        # Save outputs
        y.append(x if m.i in model.save else None)
        
        # Print layer info for 4D tensors (feature maps)
        if isinstance(x, torch.Tensor) and len(x.shape) == 4:
            B, C, H, W = x.shape
            module_name = m.__class__.__name__
            
            status = ""
            # Identify feature pyramid levels
            if H == 80 and W == 80:
                status = "← P3 CANDIDATE"
                p3_candidates.append((m.i, C, module_name))
            elif H == 40 and W == 40:
                status = "← P4 CANDIDATE"
                p4_candidates.append((m.i, C, module_name))
            elif H == 20 and W == 20:
                status = "← P5 CANDIDATE"
                p5_candidates.append((m.i, C, module_name))
            
            print(f"{m.i:<6} {module_name:<35} {str(tuple(x.shape)):<30} {H:<5} {W:<5} {C:<6}  {status}")
    
    print("\n" + "="*80)
    print("Summary - Feature Pyramid Candidates")
    print("="*80)
    
    # Display P3 candidates
    print(f"\n🔍 P3 Candidates (80x80):")
    if p3_candidates:
        for idx, channels, module in p3_candidates:
            print(f"   Layer {idx:3d}: {channels:4d} channels ({module})")
        # Recommend the last one (typically after neck processing)
        recommended_p3 = p3_candidates[-1][0]
        print(f"   ✓ RECOMMENDED: Layer {recommended_p3}")
    else:
        print("   ✗ NOT FOUND!")
    
    # Display P4 candidates
    print(f"\n🔍 P4 Candidates (40x40):")
    if p4_candidates:
        for idx, channels, module in p4_candidates:
            print(f"   Layer {idx:3d}: {channels:4d} channels ({module})")
        recommended_p4 = p4_candidates[-1][0]
        print(f"   ✓ RECOMMENDED: Layer {recommended_p4}")
    else:
        print("   ✗ NOT FOUND!")
    
    # Display P5 candidates
    print(f"\n🔍 P5 Candidates (20x20):")
    if p5_candidates:
        for idx, channels, module in p5_candidates:
            print(f"   Layer {idx:3d}: {channels:4d} channels ({module})")
        recommended_p5 = p5_candidates[-1][0]
        print(f"   ✓ RECOMMENDED: Layer {recommended_p5}")
    else:
        print("   ✗ NOT FOUND!")
    
    # Final recommendation
    if p3_candidates and p4_candidates and p5_candidates:
        p3_idx = p3_candidates[-1][0]
        p4_idx = p4_candidates[-1][0]
        p5_idx = p5_candidates[-1][0]
        
        p3_ch = p3_candidates[-1][1]
        p4_ch = p4_candidates[-1][1]
        p5_ch = p5_candidates[-1][1]
        
        print("\n" + "="*80)
        print("✅ FINAL RECOMMENDATION")
        print("="*80)
        print(f"\nUpdate extract_features_for_stata() with these indices:\n")
        print("```python")
        print(f"P3_INDEX = {p3_idx}  # [{p3_ch} channels, 80x80]")
        print(f"P4_INDEX = {p4_idx}  # [{p4_ch} channels, 40x40]")
        print(f"P5_INDEX = {p5_idx}  # [{p5_ch} channels, 20x20]")
        print("```")
        
        print(f"\nAnd update setup_stata() with channel configuration:\n")
        print("```python")
        print(f"yolo_channels = [{p3_ch}, {p4_ch}, {p5_ch}]")
        print("```")
        
        print("\n" + "="*80)
        
        return p3_idx, p4_idx, p5_idx, [p3_ch, p4_ch, p5_ch]
    else:
        print("\n" + "="*80)
        print("⚠️  WARNING: Not all feature levels found!")
        print("="*80)
        print("\nMissing:")
        if not p3_candidates:
            print("  - P3 (80x80)")
        if not p4_candidates:
            print("  - P4 (40x40)")
        if not p5_candidates:
            print("  - P5 (20x20)")
        print("\nCheck your model architecture!")
        print("="*80)
        return None, None, None, None


def main():
    """Main execution"""
    print("\n" + "="*80)
    print("YOLOv9 Layer Index Finder")
    print("="*80)
    
    # Model configuration
    model_config = 'models/detect/yolov9-c.yaml'
    
    print(f"\n[1] Loading model: {model_config}")
    try:
        model = DetectionModel(model_config, ch=3, nc=80)
        model.eval()
        print(f"    ✓ Model loaded successfully")
        print(f"    Total layers: {len(model.model)}")
    except Exception as e:
        print(f"    ✗ Failed to load model: {e}")
        return False
    
    # Find layers
    print(f"\n[2] Scanning layers for feature maps...")
    indices = find_feature_layers(model)
    
    if all(i is not None for i in indices[:3]):
        p3_idx, p4_idx, p5_idx, channels = indices
        
        # Verify by extracting actual features
        print("\n[3] Verifying indices by extracting features...")
        
        dummy_input = torch.randn(1, 3, 640, 640)
        features = []
        y = []
        
        with torch.no_grad():
            x = dummy_input
            for i, m in enumerate(model.model):
                if m.f != -1:
                    x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]
                x = m(x)
                y.append(x if m.i in model.save else None)
                
                if m.i == p3_idx:
                    features.append(x)
                    print(f"    ✓ P3 extracted: {x.shape}")
                elif m.i == p4_idx:
                    features.append(x)
                    print(f"    ✓ P4 extracted: {x.shape}")
                elif m.i == p5_idx:
                    features.append(x)
                    print(f"    ✓ P5 extracted: {x.shape}")
                
                if len(features) == 3:
                    break
        
        if len(features) == 3:
            print("\n" + "="*80)
            print("✅ SUCCESS! All feature levels verified!")
            print("="*80)
            print("\nCopy these values to your code:")
            print("\n1. In extract_features_for_stata():")
            print(f"   P3_INDEX = {p3_idx}")
            print(f"   P4_INDEX = {p4_idx}")
            print(f"   P5_INDEX = {p5_idx}")
            print("\n2. In setup_stata():")
            print(f"   yolo_channels = {channels}")
            print("\n" + "="*80 + "\n")
            return True
        else:
            print(f"\n✗ Verification failed: Expected 3 features, got {len(features)}")
            return False
    else:
        return False


if __name__ == '__main__':
    try:
        success = main()
        sys.exit(0 if success else 1)
    except Exception as e:
        print(f"\n❌ Script failed with error:")
        print(f"   {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)